use crate::{sparse::CSR, utils::global2local};

use super::sparse;
use ndarray::{concatenate, prelude::*};
use rand::prelude::*;
use rayon::{iter::{IntoParallelIterator, ParallelIterator}, result};
use std::collections::BTreeSet;

pub struct Graph<A: Clone + std::marker::Send+ std::marker::Sync> {
    adjacency: sparse::CSR<A>,
}

impl<A> Graph<A>
where
    A: Clone + std::marker::Send+ std::marker::Sync,
{
    pub fn new(
        data: ArrayView1<A>,
        row_indices: ArrayView1<usize>,
        col_indices: ArrayView1<usize>,
        num_nodes: usize,
    ) -> Graph<A> {
        let adj =
            sparse::COO::<A>::new(&data, &row_indices, &col_indices, num_nodes, num_nodes).tocsr();
        Graph::<A> { adjacency: adj }
    }

    pub fn num_node(&self) -> usize {
        self.adjacency.shape().0
    }

    pub fn num_edge(&self) -> usize {
        self.adjacency.nnz()
    }

    pub fn random_walk(
        &self,
        nodes: &ArrayView1<usize>,
        length: usize,
        num_path: usize,
    ) -> Array3<usize> {
        let mut result = Array3::<usize>::zeros((nodes.len(), num_path, length));

        for i in 0..nodes.len() {
            let mut rng = rand::thread_rng();
            for p in 0..num_path {
                let mut next_node_ind = nodes[i];
                for j in 0..length {
                    let (col_ind, _) = self.adjacency.row(next_node_ind);
                    if col_ind.len() == 0 {
                        result[[i, p, j]] = self.num_node();
                        continue;
                    }
                    next_node_ind = col_ind[rng.gen_range(0..col_ind.len())];
                    result[[i, p, j]] = next_node_ind;
                }
            }
        }
        return result;
    }
    pub fn k_hop(&self, nodes: &ArrayView1<usize>, k: usize) -> Array1<usize> {
        let mut result = BTreeSet::<usize>::new();
        let mut resultemp = BTreeSet::<usize>::new();
        let mut resulted = BTreeSet::<usize>::new();
        for &n in nodes.iter() {
            result.insert(n);
            resulted.insert(n);
            result.extend(self.adjacency.row(n).0.iter());
        }
        if k <= 1 {
            return Array1::<usize>::from_iter(result.iter().copied());
        }
        for _i in 0..k - 1 {
            for &n in result.iter() {
                if resulted.contains(&n) {
                    continue;
                } else {
                    resulted.insert(n);
                }
                resultemp.extend(self.adjacency.row(n).0.iter());
            }
            result.extend(&resultemp);
        }
        return Array1::<usize>::from_iter(result.iter().copied());
    }

    pub fn restart_random_walk(
        &self,
        nodes: &ArrayView1<usize>,
        length: usize,
        num_path: usize,
        rate: f32,
    ) -> Array3<usize> {
        // 返回概率
        let mut result = Array3::<usize>::zeros((nodes.len(), num_path, length));

        for i in 0..nodes.len() {
            let mut rng = rand::thread_rng();
            for p in 0..num_path {
                let mut next_node_ind = nodes[i];
                for j in 0..length {
                    let (col_ind, _) = self.adjacency.row(next_node_ind);
                    if col_ind.len() == 0 {
                        // result[[i, p, j]] = self.num_node();
                        next_node_ind = nodes[i];
                        continue;
                    }
                    next_node_ind = col_ind[rng.gen_range(0..col_ind.len())];
                    result[[i, p, j]] = next_node_ind;
                    if rng.gen::<f32>() < rate {
                        next_node_ind = nodes[i];
                    }
                }
            }
        }
        return result;
    }

    pub fn sample_subgraph_from_edge_by_restart_random_walk(
        &self,
        src: usize,
        dst: usize,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
    ) -> (
        Array1<usize>,
        sparse::CSR<A>,
    ) {
        let nodes = Array1::from_vec(vec![src, dst]);
        let results = self.restart_random_walk(&nodes.view(), length, num_path, rwr_rate);
        // let results = self.k_hop(&nodes.view(), 5);
        // let temp = results.len();
        // println!("{nodes}: {temp}???");
        let new_nodes = Array1::from_iter(BTreeSet::from_iter(
            results
                .iter()
                .filter(|&&x| x < self.num_node())
                .chain(nodes.iter())
                .cloned(),
        ));
        let mut local_adj = self.adjacency.submatrix(&new_nodes, &new_nodes);
        let mut nsrc = 0;
        let mut ndst = 0;
        for (i, &j) in new_nodes.iter().enumerate() {
            if j == src {
                nsrc = i;
            }
            if j == dst {
                ndst = i;
            }
        }
        local_adj.remove_element(nsrc, ndst);
        local_adj.remove_element(ndst, nsrc);
        (new_nodes, local_adj)
    }

    pub fn sample_subgraph_from_edges_by_restart_random_walk(
        &self,
        edges: &ArrayView2<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
    ) -> (
        Array1<usize>, // 子图全局节点
        Array1<usize>, // 节点对应子图index或者说edge的index，因为一个edge对应一个子图
        Array1<usize>, // batch图sender
        Array1<usize>, // batch图receiver
    ) {
        let results: Vec<(Array1<usize>, sparse::CSR<A>)> = (0..edges.shape()[0]).into_par_iter().map(|x| self.sample_subgraph_from_edge_by_restart_random_walk(edges[[x, 0]], edges[[x, 1]], length, num_path, rwr_rate)).collect();
        // let results: Vec<(Array1<usize>, sparse::CSR<A>)> = edges.rows().into_iter().map(|row| self.sample_subgraph_from_edge_by_restart_random_walk(row[0], row[1], length, num_path, rwr_rate)).collect();
        let mut nodes = Vec::new();
        let mut senders = Vec::new();
        let mut receivers = Vec::new();
        let mut labels = Vec::new();
        let mut numn = 0;
        for (i, (n, a)) in results.iter().enumerate() {
            nodes.push(n.view());
            let (_, sender, receiver) = a.tocoo().to_ndarray();
            // senders.push(n.select(Axis(0), sender.as_slice().unwrap()));
            // receivers.push(n.select(Axis(0), receiver.as_slice().unwrap()));
            let nsender = sender + numn;
            let nreceiver = receiver + numn;
            // println!("{i}: {nsender}");
            // println!("{i}: {nreceiver}");
            senders.push(nsender);
            receivers.push(nreceiver);
            labels.push(Array1::from_elem(n.len(),i));
            numn += n.len();
        };
        let mut node_array = Array1::zeros(nodes.iter().fold(0, |acc, x| acc + x.len()));
        let mut sender_array = Array1::zeros(senders.iter().fold(0, |acc,x| acc + x.len()));
        let mut receiver_array = Array1::zeros(receivers.iter().fold(0, |acc, x| acc + x.len()));
        let mut label_array = Array1::zeros(labels.iter().fold(0, |acc, x| acc + x.len()));
        let mut ns = 0;
        let mut es = 0;
        for i in 0..results.len() {
            node_array.slice_mut(s![ns..(ns+nodes[i].len())]).assign(&nodes[i]);
            sender_array.slice_mut(s![es..(es+senders[i].len())]).assign(&senders[i]);
            receiver_array.slice_mut(s![es..(es+receivers[i].len())]).assign(&receivers[i]);
            label_array.slice_mut(s![ns..(ns+labels[i].len())]).assign(&labels[i]);
            ns += nodes[i].len();
            es += senders[i].len();
        }
        return (node_array, label_array, sender_array, receiver_array);
    }

    pub fn sample_subgraph_for_bpr(&self, edges: &ArrayView2<usize>, length: usize, num_path: usize, rwr_rate: f32, num_sample: usize)  -> (
        Array1<usize>,
        Array1<usize>,
        Array1<usize>,
        Array1<usize>,
        Array2<usize>,
    ) {
        let mut rng = rand::thread_rng();
        let mut negative_edges = Array2::<usize>::zeros((edges.nrows() * num_sample * 2, 2));
        let mut ylabels = Array2::<usize>::zeros((edges.nrows(), num_sample * 2 + 1));

        for (i, edge) in edges.rows().into_iter().enumerate() {
            let (src, dst) = (edge[0], edge[1]);
            let src_neigh = self.adjacency.row(src).0.to_vec();
            let dst_neigh = self.adjacency.row(dst).0.to_vec();

            let mut src_negs = Vec::new();
            while src_negs.len() < num_sample {
                let cand = rng.gen_range(0..self.num_node());
                if !src_neigh.contains(&cand) && cand != src {
                    src_negs.push(cand);
                }
            }
            let mut dst_negs = Vec::new();
            while dst_negs.len() < num_sample {
                let cand = rng.gen_range(0..self.num_node());
                if !dst_neigh.contains(&cand) && cand != dst {
                    dst_negs.push(cand);
                }
            }
            ylabels[[i, 0]] = i;

            for j in 0..num_sample {
                ylabels[[i, j+1]] = edges.nrows() + i * num_sample * 2 + j;
                ylabels[[i, j + num_sample+1]] = edges.nrows() + i * num_sample * 2 + j + num_sample;

                negative_edges[[i * num_sample * 2 + j, 0]] = src;
                negative_edges[[i * num_sample * 2 + j, 1]] = src_negs[j];
                negative_edges[[i * num_sample * 2 + j + num_sample, 0]] = dst_negs[j];
                negative_edges[[i * num_sample * 2 + j + num_sample, 1]] = dst;
            }
        }
        let all_edges = concatenate![Axis(0), *edges, negative_edges];
        let (node, label, sender, receiver) = self.sample_subgraph_from_edges_by_restart_random_walk(&all_edges.view(), length, num_path, rwr_rate);
        (node, label, sender, receiver, ylabels)
    }

    pub fn sample_linkprediction_by_restart_random_walk(
        &self,
        nodes: &ArrayView1<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
        rate: f32,
        num_sample: usize,
    ) -> (
        Array1<usize>,
        sparse::CSR<A>,
        Array1<usize>,
        Array2<usize>,
        Array2<usize>,
    ) {
        let results = self.restart_random_walk(nodes, length, num_path, rwr_rate);
        let new_nodes = Array1::from_iter(BTreeSet::from_iter(
            results
                .iter()
                .filter(|&&x| x < self.num_node())
                .chain(nodes.iter())
                .cloned(),
        ));
        let new_adj = self.adjacency.submatrix(&new_nodes, &new_nodes);
        let local_nodes = global2local(&new_nodes.view(), nodes);
        let (a, b, c) = sample_linkprediction(new_adj, &local_nodes.view(), rate, num_sample);
        (new_nodes, a, local_nodes, b, c)
    }

    pub fn sample_by_restart_random_walk(
        &self,
        nodes: &ArrayView1<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
    ) -> (Array1<usize>, sparse::CSR<A>, Array1<usize>) {
        let results = self.restart_random_walk(nodes, length, num_path, rwr_rate);
        let new_nodes = Array1::from_iter(BTreeSet::from_iter(
            results
                .iter()
                .filter(|&&x| x < self.num_node())
                .chain(nodes.iter())
                .cloned(),
        ));
        let new_adj = self.adjacency.submatrix(&new_nodes, &new_nodes);
        let local_nodes = global2local(&new_nodes.view(), nodes);
        (new_nodes, new_adj, local_nodes)
    }
}

fn sample_linkprediction<A: Clone + std::marker::Send+ std::marker::Sync>(
    adj: sparse::CSR<A>,
    nodes: &ArrayView1<usize>,
    rate: f32,
    num_sample: usize,
) -> (sparse::CSR<A>, Array2<usize>, Array2<usize>) {
    // how many should be in train adjacency
    let mut rng = rand::thread_rng();
    let mut sender = Vec::<usize>::new();
    let mut receiver = Vec::<usize>::new();
    let mut data: Vec<A> = Vec::new();
    let mut pos_samples = Array2::zeros((nodes.shape()[0], num_sample));
    let mut neg_samples = Array2::zeros((nodes.shape()[0], num_sample));
    let nodes_vec = nodes.to_vec();
    let mut count = 0;
    for i in 0..adj.shape().0 {
        if nodes_vec.contains(&i) {
            let (neigh, weight) = adj.row(i);
            let mut neigh_vec = neigh.to_vec();
            let weight_vec = weight.to_vec();
            // let mut nneigh: Vec<usize> = (0..adj.shape().0).collect();
            // nneigh.retain(|&n| !neigh_vec.contains(&n) && n != i);
            // let nsample = nneigh
            //     .choose_multiple(&mut rng, num_sample)
            //     .cloned()
            //     .collect::<Array1<usize>>();
            let mut nsample = Vec::new();
            while nsample.len() < num_sample {
                let cand = rng.gen_range(0..adj.shape().0);
                if !neigh_vec.contains(&cand) && cand != i {
                    nsample.push(cand);
                }
            }
            neg_samples.slice_mut(s![count, ..]).assign(&Array1::from_vec(nsample));
            let isize = (neigh_vec.len() as f32 * rate).floor() as usize;
            let osize = neigh_vec.len() - isize;
            if neigh_vec.len() == 0 {
                // 不应该出现这种情况
                pos_samples.slice_mut(s![count, ..]).fill(adj.shape().0);
                neg_samples.slice_mut(s![count, ..]).fill(adj.shape().0);
                continue;
            }
            if neigh_vec.len() == 1 || osize <= 1 {
                // 较少
                pos_samples.slice_mut(s![count, ..]).fill(neigh_vec[0]);
                sender.push(i);
                receiver.push(neigh_vec[0]);
                data.push(weight_vec[0].clone());
            } else {
                let temp = (0..neigh_vec.len()).choose_multiple(&mut rng, isize);
                for &itemp in temp.iter() {
                    sender.push(i);
                    receiver.push(neigh_vec[itemp]);
                    data.push(weight_vec[itemp].clone());
                }
                let ttemp: Vec<usize> = temp.iter().map(|&n| neigh_vec[n]).collect();
                neigh_vec.retain(|&n| !ttemp.contains(&n));
                if osize > num_sample {
                    let pneigh = neigh_vec
                        .choose_multiple(&mut rng, num_sample)
                        .cloned()
                        .collect::<Array1<usize>>();
                    pos_samples.slice_mut(s![count, ..]).assign(&pneigh);
                } else {
                    let mut pneigh = Vec::<usize>::new();
                    while pneigh.len() < num_sample {
                        pneigh.push(neigh_vec[rng.gen_range(0..neigh_vec.len())]);
                    }
                    pos_samples
                        .slice_mut(s![count, ..])
                        .assign(&Array1::from_vec(pneigh));
                }
            }
            count += 1;
        } else {
            let (neigh, weight) = adj.row(i);
            for (&n, w) in neigh.iter().zip(weight.iter()) {
                sender.push(i);
                receiver.push(n);
                data.push(w.clone());
            }
        }
    }

    return (
        sparse::COO::<A>::new(
            &Array1::from_vec(data).view(),
            &Array1::from_vec(sender).view(),
            &Array1::from_vec(receiver).view(),
            adj.shape().0,
            adj.shape().1,
        )
        .tocsr(),
        pos_samples,
        neg_samples,
    );
}

#[cfg(test)]
mod tests {
    use super::*;
    use polars::prelude::*;
    use rayon::prelude::*;
    use std::error::Error;
    use std::fs::File;
    use std::path::Path;

    #[test]
    fn graph_new_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(r"/media/marcfish/backup/dataset/dataset/network/arxiv/edge.parquet");
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let g: Graph<f32> = Graph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            169343,
        );
        assert_eq!(g.num_node(), 169343);
        assert_eq!(g.num_edge(), 1166243);
        Ok(())
    }
    #[test]
    fn subgraph_k_hop_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(r"/media/marcfish/backup/dataset/dataset/network/arxiv/edge.parquet");
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let g: Graph<f32> = Graph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            169343,
        );
        assert_ne!(
            g.k_hop(&Array1::from_iter(0..1024).view(), 2).shape()[0],
            109
        );
        Ok(())
    }

    #[test]
    fn random_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(r"/media/marcfish/backup/dataset/dataset/network/arxiv/edge.parquet");
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let g: Graph<f32> = Graph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            169343,
        );
        assert_eq!(
            g.random_walk(&Array1::from_iter(0..1024).view(), 5, 5)
                .shape()[0],
            1024
        );
        Ok(())
    }

    #[test]
    fn sample_linkprediction_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(r"/media/marcfish/backup/dataset/dataset/network/arxiv/edge.parquet");
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let g: Graph<f32> = Graph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            169343,
        );
        let (a, b, c, d, e) = g.sample_linkprediction_by_restart_random_walk(
            &Array1::from_iter(0..8192).view(),
            5,
            5,
            0.3,
            0.7,
            5,
        );
        assert_ne!(a.shape()[0], 8192);
        assert_ne!(b.shape().0, 8192);
        assert_eq!(c.shape()[0], 8192);
        assert_eq!(d.shape()[0], 8192);
        assert_eq!(e.shape()[0], 8192);
        assert_ne!(d.sum(), 0);
        assert_ne!(e.sum(), 0);
        Ok(())
    }

    #[test]
    fn sample_by_restart_random_walk_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(r"/media/marcfish/backup/dataset/dataset/network/arxiv/edge.parquet");
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let g: Graph<f32> = Graph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            169343,
        );
        let (a, b, c) =
            g.sample_by_restart_random_walk(&Array1::from_iter(0..8192).view(), 5, 20, 0.3);
        assert_ne!(a.shape()[0], 8192);
        assert_ne!(b.shape().0, 8192);
        assert_eq!(c.shape()[0], 8192);
        Ok(())
    }

    #[test]
    fn sample_subgraph_for_bpr_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(r"/Users/foolfish/Desktop/project/datasets/yelp/social_5.parquet");
        let file = File::open(path)?;
        let edge = ParquetReader::new(file).finish()?.to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row_ = edge.select(Axis(1), &[0]).remove_axis(Axis(1)).mapv(|x| x as usize);
        let col_ = edge.select(Axis(1), &[1]).remove_axis(Axis(1)).mapv(|x| x as usize);
        let row = concatenate![Axis(0), row_, col_];
        let col = concatenate![Axis(0), col_, row_];
        let g: Graph<f32> = Graph::new(Array1::ones(edge.shape()[0]*2).view(), row.view(), col.view(), 319095);
        let (node, label, sender, receiver, ylabels) = g.sample_subgraph_for_bpr(&edge.slice(s![0..10000, ..]).mapv(|x| x as usize).view(), 10, 10, 0.3, 5);
        assert_eq!(ylabels.len(), 10000*11);
        let nnode = node.len();
        let nsender = sender.len();
        let nreceiver = receiver.len();
        println!("{nnode}");
        println!("{nsender}");
        println!("{nreceiver}");
        // println!("{ylabels}");
        // println!("{label}");
        // println!("{sender}");
        Ok(())
    }

    #[test]
    fn parallel_sample_linkprediction_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(r"/media/marcfish/backup/dataset/dataset/network/arxiv/edge.parquet");
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let g: Graph<f32> = Graph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            169343,
        );
        let o: Vec<(
            Array1<usize>,
            sparse::CSR<f32>,
            Array1<usize>,
            Array2<usize>,
            Array2<usize>,
        )> = (0..100)
            .into_par_iter()
            .map(|x| {
                g.sample_linkprediction_by_restart_random_walk(
                    &Array1::from_iter(x * 1024..(x + 1) * 1024).view(),
                    5,
                    5,
                    0.3,
                    0.7,
                    5,
                )
            })
            .collect();
        // let (a,b,c,d, e) = g.sample_linkprediction_by_restart_random_walk(&Array1::from_iter(0..8192), 5, 5, 0.3, 0.7, 5);
        // assert_ne!(a.shape()[0], 8192);
        // assert_ne!(b.shape().0, 8192);
        // assert_eq!(c.shape()[0], 8192);
        // assert_eq!(d.shape()[0], 8192);
        // assert_eq!(e.shape()[0], 8192);
        // assert_ne!(d.sum(), 0);
        // assert_ne!(e.sum(), 0);
        Ok(())
    }
}
